#!/usr/bin/env python3

from tkinter import ttk
import tkinter as tk
import subprocess
import time
from multiprocessing import Process, Pool, Queue
import random
import os
import signal
import sys


def process_info(num):
    print('>>> ProcessNo: {0}'.format(num))
    if hasattr(os, 'getppid'):
        print('    parent process: {0}'.format(os.getppid()))
    print('    process id: {0}'.format(os.getpid()))
    return os.getpid()

class CPU:
    def __init__(self):
        self.stopped = False
        self.pids = []
        self.q = Queue()

    def cpu_num(self):
        cpu_num = 0
        with open('/proc/cpuinfo') as f:
            for line in f:
                if line.startswith('processor'):
                    cpu_num += 1
        print('CPU NUM: {0}'.format(cpu_num))
        return cpu_num
    def start(self, cpu_num):
        for i in range(cpu_num):
            Process(target=self.operation, args=(i,self.q)).start()

    def stop(self):
        _pids = []
        self.pids += list(set([self.q.get() for _ in range(self.q.qsize())]))
        for pid in self.pids:
            r = subprocess.call(['kill', '-9', '{0}'.format(pid)])
            if r == 0:
                print('[CPU] stress stop success. pid: {0}.'.format(pid))
            else:
                _pids.append(pid)
        self.pids = _pids 


    def operation(self, num, q): 
        pid = process_info(num)
        q.put(str(pid))
        while not self.stopped:
            start = time.time()
            var = random.randint(0, 1000000000000)
            float(var) / 3.1415926

class GPU:
    def __init__(self, duration, resolution):
        self.duration = duration
        self.resolution = resolution
        self.pid = None
    def start(self):
        # glmark2 -b terrain:duration=duration -s resolution
        r = subprocess.Popen(['glmark2', '-b', 'terrain:duration={0}'.format(self.duration), '-s', self.resolution])
        self.pid = r.pid
        # subprocess.call(['glmark2', '-b', 'terrain:duration={0}'.format(self.duration), '-s', self.resolution])
    def stop(self):
        if self.pid:
            r = subprocess.call(['kill', '-9', '{}'.format(self.pid)])
            if r == 0:
                print('[GPU] stress stop success.')

        
class BenchmarkUI(object):
    def __init__(self):
        self.frame = ttk.Frame()
        self.gpu = None
        self.cpu = None
        self.padxy = {'padx': 10, 'pady': 10, 'ipady':10}
        self.resolution = tk.StringVar()


    def initUI(self, title):
        # benchmark parameter frame
        self.frame.master.title(title) 
        self.frame.pack(side='left', fill='both', expand='True')
        # monitor
        # self.jtop_btn = tk.Button(self.frame,text='Start CPU&GPU Monitor', bg='#2d8cf0', command=self.turnon_monitor)
        # self.jtop_btn.grid(row=0, column=0, **self.padxy, sticky='w')
        # benchmark
        ttk.Label(self.frame, text='Resolution').grid(row=1, column=0,pady=20)
        cb = ttk.Combobox(self.frame, textvariable=self.resolution)
        cb['values'] =('1366x768', '1920x1080', '1920x1200', '2560x1440')
        cb.grid(row=1, column=0, sticky='e')
        cb.current(1)

        self.start_btn = tk.Button(self.frame, text='Start Test', bg='#2d8cf0', command=self.start_benchmark)
        self.start_btn.grid(row=1, column=1, sticky='w', padx=20)
        self.stop_btn = tk.Button(self.frame, text='Stop Test', state='disabled', command=self.stop_benchmark)
        self.stop_btn.grid(row=1, column=1)
       
        termf = tk.Frame(self.frame, width=800, height=600)
        termf.grid(row=4, columnspan=2, padx=10, pady=10)
        wid = termf.winfo_id()
        subprocess.Popen(['xterm', '-into', f'{wid}', '-geometry', '800x600', '-e', 'jtop'])

    def start_benchmark(self):
        print('sart stress...') 
        self.stop_btn['state'] = 'normal'
        self.start_btn['state'] = 'disabled'
        resolution_val = self.resolution.get().strip()
        print(resolution_val)
        
        if not resolution_val:
            resolution_val = '2560x1440'
            self.resolution_input.set('2560x1440')
        # default gpu test duration 
        duration_val = 86400 

        # CPU
        self.cpu = CPU()
        cpu_val = self.cpu.cpu_num()
        self.cpu.start(cpu_val)

        # GPU
        self.gpu = GPU(duration_val, resolution_val)
        self.gpu.start()
    
    def stop_benchmark(self):
        print('stop stress') 
        self.gpu.stop()
        self.cpu.stop()
        self.stop_btn['state'] = 'disabled'
        self.start_btn['state'] = 'normal'
    
    def turnon_monitor(self):
        # subprocess.Popen(['jtop'])
        subprocess.Popen(['gnome-terminal','--title=Monitor', '--', 'jtop'])
    
def run():
    root = tk.Tk()
    # root.geometry('310x160+700+300')
    root.geometry('820x700+700+300')
    frame = BenchmarkUI()
    frame.initUI('GPU&CPU Stress Tool')
    root.mainloop()


if __name__ == '__main__':
    run()


